package com.ververica.cdc.guass.sink.jdbc.internal.executor;

import com.ververica.cdc.guass.sink.jdbc.core.database.dialect.JdbcDialect;
import com.ververica.cdc.guass.sink.jdbc.core.database.dialect.JdbcDialectConverter;
import com.ververica.cdc.guass.sink.jdbc.internal.JdbcOutputFormat;
import com.ververica.cdc.guass.sink.jdbc.statement.FieldNamedPreparedStatement;
import com.ververica.cdc.guass.sink.jdbc.util.JsonUtil;
import com.ververica.cdc.guass.source.kafka.data.NullableValue;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.types.RowKind;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.*;
import java.util.*;
import java.util.function.Function;

public class TableInsertOrUpdateDynamicFieldsStatementExecutor implements JdbcBatchStatementExecutor<RowData> {

    private final JdbcDialect dialect;
    private final JdbcDialectConverter existSetter;
    private final String tableName;
    private final String[] fieldNames;
    private final LogicalType[] fieldTypes;
    private final String[] pkNames;

    private Connection connection;

    // Cache structure: key is the list of fields to update, value is the list of rows
    private final Map<List<String>, List<RowData>> updateRows = new HashMap<>();
    private final Map<List<String>, List<RowData>> insertRows = new HashMap<>();

    private String existStmt;
    private FieldNamedPreparedStatement existPreparedStatement;

    private final Map<List<String>, PreparedStatementFactory> insertStmtFactories = new HashMap<>();
    private final Map<List<String>, PreparedStatementFactory> updateStmtFactories = new HashMap<>();

    private final Function<RowData, RowData> keyExtractor;

    private static final Logger LOG = LoggerFactory.getLogger(TableInsertOrUpdateDynamicFieldsStatementExecutor.class);

    // 新增当前批次ID跟踪
    private String currentBatchId;

    /**
     * Functional interface for creating PreparedStatements.
     */
    @FunctionalInterface
    interface PreparedStatementFactory {
        FieldNamedPreparedStatement create(Connection connection) throws SQLException;
    }

    public TableInsertOrUpdateDynamicFieldsStatementExecutor(
            JdbcDialect dialect,
            JdbcDialectConverter existSetter,
            String tableName,
            String[] fieldNames,
            LogicalType[] fieldTypes,
            String[] pkNames,
            Function<RowData, RowData> keyExtractor) {
        this.dialect = dialect;
        this.existSetter = existSetter;
        this.fieldNames = fieldNames;
        this.fieldTypes = fieldTypes;
        this.tableName = tableName;
        this.pkNames = pkNames;
        this.keyExtractor = keyExtractor;
    }

    @Override
    public void prepareStatements(Connection connection) throws SQLException {
        this.connection = connection;
        this.existStmt = dialect.getRowExistsStatement(tableName, pkNames);
        this.existPreparedStatement = FieldNamedPreparedStatement.prepareStatement(connection, existStmt, pkNames);
        LOG.info("Prepared statements for table {}: Exist Statement prepared.", tableName);
    }

    @Override
    public void addToBatch(RowData record) throws SQLException {

        // 确保批次ID初始化
        if (currentBatchId == null) {
            currentBatchId = UUID.randomUUID().toString();
            LOG.info("Started new batch {}", currentBatchId);
        }

        LOG.debug("Adding record to batch: {}", record);
        RowData pk = keyExtractor.apply(record);
        RowData convertRowData = convertRealValue(record);

        if (exist(pk)) {
            // Extract fields present in this record excluding primary keys
            List<String> presentFields = getPresentFields(record, fieldNames);
            if (presentFields.isEmpty()) {
                // Nothing to update
                LOG.info("No fields to update in record: {}", record);
                return;
            }
            updateRows.computeIfAbsent(presentFields, k -> new ArrayList<>()).add(convertRowData);
            LOG.debug("Added record to update batch for fields: {}", presentFields);
        } else {

            if(record.getRowKind().equals(RowKind.UPDATE_AFTER)){
                throw new RuntimeException( "record "+ JsonUtil.toJson(record) +" prime key not exists can't update");
            }

            // Handle insert
            List<String> presentFields = getPresentFields(record, fieldNames);
            if (presentFields.isEmpty()) {
                // No fields to insert
                LOG.info("No fields to insert in record: {}", record);
                return;
            }
            insertRows.computeIfAbsent(presentFields, k -> new ArrayList<>()).add(convertRowData);
            LOG.debug("Added record to insert batch for fields: {}", presentFields);
        }
    }

    private boolean exist(RowData pk) throws SQLException {
        LOG.debug("Checking if record exists with primary key: {}", pk);
        existSetter.toExternal(pk, existPreparedStatement);
        try (ResultSet resultSet = existPreparedStatement.executeQuery()) {
            boolean exists = resultSet.next();
            LOG.debug("Exist check result for primary key {}: {}", pk, exists);
            return exists;
        }
    }

    @Override
    public void executeBatch() throws SQLException {

        LOG.info("Executing batch insert or update.");

        // Handle updates
        for (Map.Entry<List<String>, List<RowData>> entry : updateRows.entrySet()) {
            List<String> fieldsToUpdate = entry.getKey();
            List<RowData> rows = entry.getValue();
            PreparedStatementFactory updateFactory = updateStmtFactories.get(fieldsToUpdate);
            if (updateFactory == null) {
                String[] fieldArray = fieldsToUpdate.toArray(new String[0]);
                String updateStmt = dialect.getUpdateStatement(tableName, fieldArray, pkNames);
                LOG.info("Executing batch with statement: {} batchId {}", updateStmt,currentBatchId);
                updateFactory = conn -> FieldNamedPreparedStatement.prepareStatement(conn, updateStmt, fieldsToUpdate.toArray(new String[0]));
                updateStmtFactories.put(fieldsToUpdate, updateFactory);
            }
            LOG.debug("Executing update batch for fields: {} batchId {}", fieldsToUpdate,currentBatchId);
            executeBatchWithErrorHandling(updateFactory, rows, createPartialConverter(fieldsToUpdate));
        }

        // Handle inserts
        for (Map.Entry<List<String>, List<RowData>> entry : insertRows.entrySet()) {
            List<String> fieldsToInsert = entry.getKey();
            List<RowData> rows = entry.getValue();
            PreparedStatementFactory insertFactory = insertStmtFactories.get(fieldsToInsert);
            if (insertFactory == null) {
                String[] fieldArray = fieldsToInsert.toArray(new String[0]);
                String insertStmt = dialect.getInsertIntoStatement(tableName, fieldArray);
                LOG.info("Executing batch with statement: {} batchId {}", insertStmt,currentBatchId);
                insertFactory = conn -> FieldNamedPreparedStatement.prepareStatement(conn, insertStmt, fieldArray);
                insertStmtFactories.put(fieldsToInsert, insertFactory);
            }
            LOG.debug("Executing insert batch for fields: {} batchId {}", fieldsToInsert,currentBatchId);
            executeBatchWithErrorHandling(insertFactory, rows, createPartialConverter(fieldsToInsert));
        }

        // Clear caches after execution
        updateRows.clear();
        insertRows.clear();
        currentBatchId = null; // 重置批次ID
        LOG.info("Batch {} execution completed and caches cleared.",currentBatchId);
    }

    private void executeBatchWithErrorHandling(PreparedStatementFactory statementFactory, List<RowData> rows, JdbcDialectConverter converter) throws SQLException {

        // 打印调试信息：开始执行批处理
        LOG.info("batchId:{} Number of rows:{} ",currentBatchId, rows.size());

        FieldNamedPreparedStatement ps = null;
        try  {
            ps = statementFactory.create(connection);
            for (RowData row : rows) {
                converter.toExternal(row, ps);
                ps.addBatch();
            }
            try {
                ps.executeBatch();
                LOG.info("Batch {} executed successfully.",currentBatchId);
            } catch (BatchUpdateException e) {
                handleBatchUpdateException(e, rows);
            }
        }catch (SQLException e) {
            LOG.error("Error executing batch.", e);
            throw e;
        }
    }


    private void handleBatchUpdateException(BatchUpdateException e, List<RowData> rows) {
        int[] updateCounts = e.getUpdateCounts();
        for (int i = 0; i < updateCounts.length; i++) {
            if (updateCounts[i] == Statement.EXECUTE_FAILED) {
                handleFailedStatement(rows.get(i));
            }
        }
    }

    private void handleFailedStatement(RowData row) {
        LOG.error("Failed to execute statement for row: {}, batchId: {} ", row,currentBatchId);
    }


    @Override
    public void closeStatements() throws SQLException {
        if (existPreparedStatement != null) {
            existPreparedStatement.close();
            LOG.info("Closed existPreparedStatement for existence check.");
        }

    }


    /**
     * 将包含 NullableValue 的 RowData 转换为仅包含实际值或 null 的 RowData。
     *
     * @param record 原始的 RowData 对象。
     * @return 转换后的 RowData 对象，所有字段为实际值或 null。
     */
    private RowData convertRealValue(RowData record) {
        if (record == null) {
            return null;
        }

        int validFieldCount = 0;
        RowData newRowData = null;

        if (record instanceof GenericRowData) {

            GenericRowData genericRowData = (GenericRowData) record;
            for (int i = 0; i < record.getArity(); i++) {
                NullableValue<?> nullable = (NullableValue<?>) genericRowData.getField(i);
                if (nullable.isExists()) {
                    validFieldCount++;
                }
            }
            newRowData = new GenericRowData(record.getRowKind(), validFieldCount);
            int j = 0;
            for (int i = 0; i < record.getArity(); i++) {

                NullableValue nullable = (NullableValue) ((GenericRowData) record).getField(i);
                if (nullable.isExists()) {
                    ((GenericRowData) newRowData).setField(j, nullable.getValue());
                    j++;
                }
            }
        } else {
            newRowData = record;
        }


        return newRowData;
    }


    /**
     * Extracts the fields that are present (non-null) in the RowData, excluding primary keys.
     *
     * @param rowData    The row data.
     * @param fieldNames The field names to check.
     * @return A list of field names that are present.
     */
    private List<String> getPresentFields(RowData rowData, String[] fieldNames) {

        List<String> presentFields = new ArrayList<>();
        for (int i = 0; i < fieldNames.length; i++) {
            String field = fieldNames[i];
            boolean fieldExists = false;
            if (rowData instanceof GenericRowData) {
                NullableValue nv = (NullableValue) ((GenericRowData) rowData).getField(i);
                fieldExists = nv.isExists();
            } else {
                fieldExists = !rowData.isNullAt(i);
            }

            if (fieldExists) {
                presentFields.add(field);
            }

        }
        return presentFields;
    }


    /**
     * 创建一个仅处理指定字段集合的converter。
     */
    private JdbcDialectConverter createPartialConverter(List<String> existFieldNames) {
        LogicalType[] existLogicalTypes = new LogicalType[existFieldNames.size()];
        for (int i = 0; i < existFieldNames.size(); i++) {
            int idx = Arrays.asList(fieldNames).indexOf(existFieldNames.get(i));
            if (idx == -1) {
                throw new IllegalArgumentException("Field '" + existFieldNames.get(i) + "' not found in allFieldNames");
            }
            existLogicalTypes[i] = fieldTypes[idx];
        }
        RowType rowType = RowType.of(existLogicalTypes, existFieldNames.toArray(new String[0]));
        return dialect.getRowConverter(rowType);
    }

}
